package net.lab1024.sa.gateway.filter;

import net.lab1024.sa.gateway.util.TrafficUtils;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * 限流过滤器
 */
@Component
public class RateLimiterFilter implements GatewayFilter {
    private final RedisRateLimiter redisRateLimiter;

    public RateLimiterFilter(RedisRateLimiter redisRateLimiter) {
        this.redisRateLimiter = redisRateLimiter;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        String path = exchange.getRequest().getPath().value();
        String limitKey = TrafficUtils.isMultipartRequest(exchange) ? "file_upload:" + TrafficUtils.getClientIp(exchange) : TrafficUtils.getClientIp(exchange);
        String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
        return redisRateLimiter.isAllowed(ip, String.valueOf(1)) // 按IP限流
                .flatMap(limit -> {
                    if (!limit.isAllowed()) {
                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                        return exchange.getResponse().writeWith(Mono.just(
                                exchange.getResponse().bufferFactory().wrap("Rate limit exceeded".getBytes())
                        ));
                    }
                    return chain.filter(exchange); // 放行
                });
    }

}